import ollama
from ollama import ChatResponse
import logging

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s')
logger = logging.getLogger(__name__)

class Conversation:
    def __init__(self, prompt, num_of_round):
        self.prompt = prompt
        self.num_of_round = num_of_round
        self.messages = []
        self.messages.append({"role": "system", "content": self.prompt})

    def ask(self, question):
        try:
            self.messages.append({"role": "user", "content": question})
            response = ollama.chat(
                model="qwen3:0.6b",
                messages=self.messages,
            )
        except Exception as e:
            print(e)
            return e
        
        #message = response["choices"][0]["message"]["content"]
        message = response.message.content
        #logging.info(response)
        self.metrics(response)
        self.messages.append({"role": "assistant", "content": message})

        if len(self.messages) > self.num_of_round * 2 + 1:
            del self.messages[1:3]  # Remove the first round conversation left.
        return message
    def metrics(self, response: ChatResponse):
        metrics={
            "created_at": response.created_at,
            "model": response.model,
            "prompt":{
                "token": response.prompt_eval_count,
                "rate": round(response.prompt_eval_count/response.prompt_eval_duration*10**9,2),
            },
            "token": response.eval_count,
            "rate": round(response.eval_count/response.eval_duration*10**9,2),
            "total": response.prompt_eval_count+response.eval_count,
        }
        logger.info(f"ollama metrics: {metrics}")

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    prompt = """你是一个中国厨师，用中文回答做菜的问题。你的回答需要满足以下要求:
    1. 你的回答必须是中文
    2. 回答限制在100个字以内"""
    conv1 = Conversation(prompt, 3)

    question1 = "你是谁？"
    print("User : %s" % question1)
    print("Assistant : %s\n" % conv1.ask(question1))

    question2 = "请问鱼香肉丝怎么做？"
    print("User : %s" % question2)
    print("Assistant : %s\n" % conv1.ask(question2))

    question3 = "那蚝油牛肉呢？"
    print("User : %s" % question3)
    print("Assistant : %s\n" % conv1.ask(question3))

    question4 = "我问你的第一个问题是什么？"
    print("User : %s" % question4)
    print("Assistant : %s\n" % conv1.ask(question4))
